from inspect import cleandoc
import io
import os
import time
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from PIL import Image
from io import BytesIO
import numpy as np
import av
import torchaudio

def create_session_with_retries(max_retries=3, backoff_factor=0.3):
    """
    创建一个带有重试机制的requests会话
    
    Args:
        max_retries: 最大重试次数
        backoff_factor: 重试间隔时间因子
    
    Returns:
        requests.Session: 配置好重试机制的会话对象
    """
    session = requests.Session()
    retry_strategy = Retry(
        total=max_retries,
        backoff_factor=backoff_factor,
        status_forcelist=[429, 500, 502, 503, 504],  # 需要重试的HTTP状态码
        allowed_methods=["GET", "POST"],  # 允许重试的HTTP方法
    )
    adapter = HTTPAdapter(max_retries=retry_strategy)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    # 设置超时时间
    session.timeout = 30  # 30秒超时
    return session

# 创建全局会话对象
_session = create_session_with_retries()


def save_to_url(file_content, url, filename=None):
    """
    将文件内容上传到指定的URL，带有重试机制
    
    Args:
        file_content: 文件内容，可以是BytesIO对象或字节数据
        url: 目标URL
        filename: 可选的文件名
    
    Returns:
        str: 上传成功后的URL或响应信息
    
    Raises:
        Exception: 当上传失败时抛出异常
    """
    # 准备文件内容
    if hasattr(file_content, 'read'):
        file_content.seek(0)  # 确保从头开始读取
        # 保存内容到临时变量，以便在重试时可以重新读取
        content_data = file_content.read()
        files = {'file': (filename or 'uploaded_file', content_data)}
    else:
        content_data = file_content
        files = {'file': (filename or 'uploaded_file', content_data)}
    
    # 发送POST请求，使用带重试机制的会话
    max_retries = 3
    retry_count = 0
    
    while retry_count <= max_retries:
        try:
            print(f"Attempting to upload to {url} (attempt {retry_count + 1}/{max_retries + 1})...")
            response = _session.post(url, files=files, timeout=30)
            response.raise_for_status()  # 检查是否有HTTP错误
            
            # 尝试解析JSON响应，如果失败则返回文本响应
            try:
                result = response.json()
                # 尝试从响应中提取URL，如果没有则返回响应内容
                if isinstance(result, dict):
                    return result.get('url', result.get('filename', str(result)))
                else:
                    return str(result)
            except ValueError:
                return response.text or url
                
        except requests.exceptions.RequestException as e:
            retry_count += 1
            error_message = f"Error uploading to URL {url}: {str(e)}"
            print(error_message)
            
            if retry_count <= max_retries:
                # 计算退避时间
                backoff_time = 2 ** (retry_count - 1)  # 指数退避
                print(f"Retrying in {backoff_time} seconds...")
                time.sleep(backoff_time)
            else:
                # 最后一次重试也失败了
                raise Exception(f"Failed to upload after {max_retries + 1} attempts: {error_message}") from e
        except Exception as e:
            print(f"Unexpected error uploading to URL {url}: {str(e)}")
            raise e


def save_to_local(file_content, filename):
    """
    简单的本地保存函数，替代Azure Blob存储
    """
    # 确保保存目录存在
    save_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'saved_media')
    os.makedirs(save_dir, exist_ok=True)
    
    # 保存文件到本地
    file_path = os.path.join(save_dir, filename)
    with open(file_path, 'wb') as f:
        # 如果是BytesIO对象，需要读取其内容
        if hasattr(file_content, 'read'):
            file_content.seek(0)
            f.write(file_content.read())
        else:
            f.write(file_content)
    
    return file_path


class SaveImageToRemote:
    """
    一个将图像保存到远程URL或本地存储的节点。
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE", {"tooltip": "The image to save."}),
                "filename": (
                    "STRING",
                    {"tooltip": "The filename to save."},
                ),
            },
            "optional": {
                "remote_url": (
                    "STRING",
                    {"tooltip": "URL to upload the file to. If empty, saves locally."},
                ),
            },
        }

    RETURN_TYPES = ()
    DESCRIPTION = cleandoc(__doc__)
    FUNCTION = "save_image"

    OUTPUT_NODE = True
    OUTPUT_TOOLTIPS = ("Saved image path or URL",)  # Tooltips for the output node

    CATEGORY = "image"

    def tensor_to_pil(self, image_tensor):
        """
        将torch张量转换回PIL图像列表。
        """
        # Ensure 4D shape: (frames, H, W, 3)
        if image_tensor.ndim == 3:
            image_tensor = image_tensor.unsqueeze(0)

        frames = []
        for i in range(image_tensor.shape[0]):
            rgb = image_tensor[i].clamp(0, 1).numpy()  # (H, W, 3)
            rgb_uint8 = (rgb * 255).astype(np.uint8)
            frame = Image.fromarray(rgb_uint8, mode="RGB")
            frames.append(frame)

        return frames

    def save_image(self, image, filename, remote_url=""):
        format = filename.split(".")[-1].lower()

        if format not in ["png"]:
            raise ValueError("SaveImageToRemote only supports PNG format.")

        frames = self.tensor_to_pil(image)

        file = BytesIO()
        frames[0].save(file, format=format)
        file.seek(0)

        # 保存到远程URL或本地
        try:
            if remote_url:
                # 上传到指定URL
                result = save_to_url(file, remote_url, filename)
                print(f"Image uploaded to URL: {result}")
                saved_location = result
            else:
                # 保存到本地
                saved_location = save_to_local(file, filename)
                print(f"Image saved locally to: {saved_location}")
        except Exception as e:
            print("Error saving image:", e)
            raise e

        return {"ui": {"urls": (saved_location,)}}


class SaveVideoToRemote:
    """
    一个将视频保存到远程URL或本地存储的节点。
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "video": ("VIDEO", {"tooltip": "The video to save."}),
                "filename": (
                    "STRING",
                    {"tooltip": "The filename to save."},
                ),
            },
            "optional": {
                "remote_url": (
                    "STRING",
                    {"tooltip": "URL to upload the file to. If empty, saves locally."},
                ),
            },
        }

    RETURN_TYPES = ()
    DESCRIPTION = cleandoc(__doc__)
    FUNCTION = "save_video"
    CATEGORY = "image/video"

    OUTPUT_NODE = True
    OUTPUT_TOOLTIPS = ("Saved video path or URL",)  # Tooltips for the output node

    def save_video(self, video, filename, remote_url=""):
        try:
            # 这里需要根据ComfyUI的video对象结构进行调整
            # 先保存到临时文件
            temp_path = os.path.join(os.getenv('TEMP', '/tmp'), filename)
            
            # 假设video对象有save_to方法
            if hasattr(video, 'save_to'):
                video.save_to(temp_path)
            else:
                # 占位：如果没有save_to方法，这里需要实现其他保存逻辑
                raise NotImplementedError("Video saving not fully implemented")
            
            # 读取临时文件内容
            with open(temp_path, 'rb') as f:
                video_content = f.read()
            
            # 保存到远程URL或本地
            if remote_url:
                result = save_to_url(video_content, remote_url, filename)
                print(f"Video uploaded to URL: {result}")
                saved_location = result
            else:
                saved_location = save_to_local(video_content, filename)
                print(f"Video saved locally to: {saved_location}")
            
            # 清理临时文件
            if os.path.exists(temp_path):
                os.remove(temp_path)
                
            return {"ui": {"urls": (saved_location,)}}
        except Exception as e:
            print("Error saving video:", e)
            raise e


class SaveAudioToRemote:
    """
    一个将音频保存到远程URL或本地存储的节点。
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "audio": ("AUDIO", {"tooltip": "The audio to save."}),
                "filename": (
                    "STRING",
                    {"tooltip": "The filename to save."},
                ),
            },
            "optional": {
                "remote_url": (
                    "STRING",
                    {"tooltip": "URL to upload the file to. If empty, saves locally."},
                ),
            },
        }

    RETURN_TYPES = ()
    DESCRIPTION = cleandoc(__doc__)
    FUNCTION = "save_audio"
    CATEGORY = "audio"

    OUTPUT_NODE = True
    OUTPUT_TOOLTIPS = ("Saved audio path or URL",)  # Tooltips for the output node

    def save_audio(self, audio, filename, remote_url=""):
        try:
            # 创建一个BytesIO对象来存储音频数据
            output_buffer = io.BytesIO()
            
            # 简化的音频处理逻辑，实际实现需要根据audio对象的结构进行调整
            # 这里假设audio是一个字典，包含waveform和sample_rate
            if isinstance(audio, dict) and 'waveform' in audio and 'sample_rate' in audio:
                # 使用torchaudio保存音频
                # 注意：这需要根据实际的audio对象结构进行调整
                # 这里只是一个示例实现
                torchaudio.save(output_buffer, audio['waveform'], audio['sample_rate'], format=filename.split('.')[-1].lower())
                output_buffer.seek(0)
                
                # 保存到远程URL或本地
                if remote_url:
                    result = save_to_url(output_buffer, remote_url, filename)
                    print(f"Audio uploaded to URL: {result}")
                    saved_location = result
                else:
                    saved_location = save_to_local(output_buffer, filename)
                    print(f"Audio saved locally to: {saved_location}")
                    
                return {"ui": {"urls": (saved_location,)}}
            else:
                raise ValueError("Invalid audio object format")
        except Exception as e:
            print("Error saving audio:", e)
            raise e
